import pandas as pd
import joblib
from utils.common import preprocess_data, plot_auc_curve
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score


def predict(data_path, model_dir, results_dir, output_path=None):
    """
    使用训练好的模型进行预测[3](@ref)
    参数:
        data_path: 预测数据路径
        model_dir: 模型目录
        results_dir: 结果目录
        output_path: 预测结果保存路径
    """
    # 1. 加载模型和预处理对象
    model = joblib.load(f'{model_dir}/xgb_model.pkl')
    scaler = joblib.load(f'{model_dir}/scaler.pkl')
    feature_columns = joblib.load(f'{model_dir}/feature_columns.pkl')

    # 2. 加载并预处理数据
    new_data = pd.read_csv(data_path)
    processed = preprocess_data(new_data, feature_columns)

    # 3. 安全提取特征（兼容有无目标列）
    if 'Attrition' in processed.columns:
        X_new = processed.drop(columns=['Attrition'])
        y_true = processed['Attrition']
    else:
        X_new = processed.copy()
        y_true = None

    # 4. 标准化和预测
    X_new_scaled = scaler.transform(X_new)
    predictions = model.predict_proba(X_new_scaled)[:, 1]  # 获取正类概率

    # 5. 如果有真实标签，评估并绘制AUC曲线
    if y_true is not None:
        pred_auc = roc_auc_score(y_true, predictions)
        print(f" 预测集AUC分数: {pred_auc:.4f}")
        plot_auc_curve(
            y_true, predictions,
            dataset_name="预测集",
            save_path=f'{results_dir}/pred_auc_curve.png'
        )

    # 6. 保存结果
    if output_path:
        result_df = pd.DataFrame({
            'EmployeeID': new_data['EmployeeNumber'],
            'Attrition_Probability': predictions
        })


    return predictions


if __name__ == "__main__":
    # 配置路径
    data_path = '../data/test2.csv'  # 实际预测数据
    model_dir = '../model'
    results_dir = '../data/fig'


    # 执行预测
    predictions = predict(data_path, model_dir, results_dir)
    print(f" 生成 {len(predictions)} 条预测结果")
    print(" 前5条预测概率:")
    for i, prob in enumerate(predictions[:5]):
        print(f"  样本{i + 1}: {prob:.4f}")